from typing import TypedDict

from langgraph.constants import START, END
from langgraph.graph import StateGraph


class InputState(TypedDict):
    user_input: str


class OutputState(TypedDict):
    graph_output: str


class OverallState(TypedDict):
    foo: str
    user_input: str
    graph_output: str


class PrivateState(TypedDict):
    bar: str


def node_1(state: InputState) -> OverallState:
    return {"foo": state["user_input"] + " & lcc"}


def node_2(state: OverallState) -> PrivateState:
    return {"bar": state["foo"] + " sweet"}


def node_3(state: PrivateState) -> OutputState:
    return {"graph_output": state["bar"] + " & happy"}


# 构建图
builder = StateGraph(
    state_schema=OverallState,
    input_schema=InputState,
    output_schema=OutputState
)

# 添加node
builder.add_node("node_1", node_1)
builder.add_node("node_2", node_2)
builder.add_node("node_3", node_3)

# 添加edge
builder.add_edge(START, "node_1")
builder.add_edge("node_1", "node_2")
builder.add_edge("node_2", "node_3")
builder.add_edge("node_3", END)

# 编译图
graph = builder.compile()

# 调用图
content = graph.invoke({"user_input": "lh"})
print(content)

# 输出mermaid
print(graph.get_graph().draw_mermaid())
